from collections import OrderedDict
import copy
import torch
from models.resnet_hetero import ResNet18_cifar,ResNet18_widar
from models.vgg_hetero import vgg_16_bn
from models.ResNet_Depth import multi_resnet18
import numpy as np

def Aggregation(w,lens):
    w_avg = None
    if len(w)==0:
        return w_avg
    total_count = sum(lens)

    for i in range(0,len(w)):
        if i== 0:
            w_avg = copy.deepcopy(w[0])
            for k in w_avg.keys():
                w_avg[k] = w[i][k]*lens[i]
        else:
            for k in w_avg.keys():
                w_avg[k] += w[i][k]*lens[i]

    for k in w_avg.keys():
        w_avg[k] = torch.div(w_avg[k],total_count)
    return w_avg
 
def AggregationNoData(w):
    w_avg = None
    if len(w)==0:
        return w_avg
    for i in range(0, len(w)):
        if i == 0:
            w_avg = copy.deepcopy(w[0])
            for k in w_avg.keys():
                w_avg[k] = copy.deepcopy(w[i][k])
        else:
            for k in w_avg.keys():
                w_avg[k] += w[i][k]

    for k in w_avg.keys():
        w_avg[k] = torch.div(w_avg[k], len(w))

    return w_avg

def split_model(global_param, slim_param):
    param = copy.deepcopy(slim_param)
    for k, v in param.items():  # 遍历所有层，每一层遍历
        if v.dim() > 1:
            d1 = v.shape[0]
            d2 = v.shape[1]
            param[k] = global_param[k][:d1, :d2]
        elif v.dim() == 1:
            d1 = v.shape[0]
            param[k] = global_param[k][:d1]
        else:
            param[k] = global_param[k]
    return param


def Aggregation_FedSlim(w, lens, global_model_param):
    w_avg = copy.deepcopy(global_model_param)  # largest model
    count = OrderedDict()
    for k, v in w_avg.items():  # 遍历所有层，每一层遍历
        parameter_type = k.split('.')[-1]

        count[k] = v.new_zeros(v.size(), dtype=torch.float32)
        tmp_v = v.new_zeros(v.size(), dtype=torch.float32)
        for m in range(len(w)):  # 遍历所有用户
            if parameter_type == 'weight':
                if v.dim() > 1:  # 卷积  和 线性层
                    d1 = w[m][k].shape[0]
                    d2 = w[m][k].shape[1]
                    tmp_v[:d1, :d2] += w[m][k] * lens[m]  # 第m个客户端的 k 层参数
                    count[k][:d1, :d2] += lens[m]
                else:  # BN层
                    d1 = w[m][k].shape[0]
                    tmp_v[:d1] += w[m][k] * lens[m]
                    count[k][:d1] += lens[m]
            else:
                d1 = w[m][k].shape[0]
                tmp_v[:d1] += w[m][k] * lens[m]
                count[k][:d1] += lens[m]

        tmp_v[count[k] > 0] = tmp_v[count[k] > 0].div_(count[k][count[k] > 0])
        tmp_v[count[k] == 0] = global_model_param[k][count[k] == 0]
        w_avg[k] = tmp_v

    return w_avg


def get_model_list(args):
    model_rate = args.width_ration
    depth_list = args.depth_saved

    net_glob_list = []
    net_slim_info = []
    for i in model_rate:
        for depth in depth_list:
            if args.model == 'resnet18':
                if args.dataset == 'widar':
                    net = ResNet18_widar(num_channels=args.num_channels,num_classes=args.num_classes, track_running_stats=False, slim_idx=depth, scale=i)
                else:
                    net = ResNet18_cifar(num_channels=args.num_channels, num_classes=args.num_classes, track_running_stats=False, slim_idx=depth, scale=i)
            elif args.model == 'vgg16':
                net = vgg_16_bn(num_classes=args.num_classes, track_running_stats=False, num_channels=args.num_channels, slim_idx=depth, scale=i)
            else:
                print("model error")

            total = sum([param.nelement() for param in net.parameters()])
            net.to(args.device)
            net.train()
            print("==" * 50)
            print('【model config】  model_name:{}, width:{} , depth:{}, param:{}MB'.format(args.model, i, depth, total * 4 / 1e6))
            print(net)
            net_glob_list.append(net)
            net_slim_info.append((i, depth, total / 1e6))  # 宽度 深度 参数量

            if i == 1.0:
                break
    return net_glob_list, net_slim_info

def get_model_list_depthfl(args):
    n_blocks = args.n_blocks

    net_glob_list = []
    net_slim_info = []
    if args.model == 'resnet18':
        for depth in n_blocks:
            net = multi_resnet18(n_blocks = depth, num_classes=args.num_classes)
            total = sum([param.nelement() for param in net.parameters()])
            net.to(args.device)
            net.train()
            print("==" * 50)
            print('【model config】  model_name:{}, n_blocks:{}, param:{}M'.format(args.model, depth,
                                                                                    total / 1e6))
            # print(net)
            net_glob_list.append(net)
            net_slim_info.append((depth, total / 1e6))  # 宽度 深度 参数量


    else:
        print("model error")
    return net_glob_list,net_glob_list

def select_clients(args ):
    m = max(int(args.frac*args.num_users),1)
    m1 = int(m * args.client1_frac)
    m2 = int(m * args.client2_frac)
    m3 = int(m * args.client3_frac)
    idxs_users = []
    ration_users = []
    # idxs_users.extend(np.random.choice(range(0, int(args.num_users)), m1, replace=False))
    # idxs_users.extend(
    #     np.random.choice(range(int(args.num_users * args.client1_frac), int(args.num_users)), m2, replace=False))
    # idxs_users.extend(
    #     np.random.choice(range(int(args.num_users * (args.client1_frac + args.client2_frac)), int(args.num_users)), m3,
    #                      replace=False))
    l1 = range(0, int(args.num_users))
    l2 = range(int(args.num_users * args.client1_frac), int(args.num_users))
    l3 = range(int(args.num_users * (args.client1_frac + args.client2_frac)), int(args.num_users))

    c3 = np.random.choice(l3, m3, replace=False)
    l2 = [x for x in l2 if x not in c3]
    l1 = [x for x in l1 if x not in c3]

    c2 = np.random.choice(l2, m2, replace=False)
    l1 = [x for x in l1 if x not in c2]
    c1 = np.random.choice(l1, m1, replace=False)
    idxs_users.extend(c1)
    idxs_users.extend(c2)
    idxs_users.extend(c3)
    for _ in range(m1):
        ration_users.append(0)
    for _ in range(m2):
        ration_users.append(1)
    for _ in range(m3):
        ration_users.append(2)
    return ration_users,idxs_users

def select_clients_s(args ):
    m = max(int(args.frac*args.num_users),1)
    m1 = int(m * args.client1_frac)
    m2 = int(m * args.client2_frac)
    m3 = int(m * args.client3_frac)
    idxs_users = []
    ration_users = []
    l1 = range(0, int(args.num_users))
    l2 = range(int(args.num_users * args.client1_frac), int(args.num_users))
    l3 = range(int(args.num_users * (args.client1_frac + args.client2_frac)), int(args.num_users))

    c3 = np.random.choice(l3, m3, replace=False)
    l2 = [x for x in l2 if x not in c3]
    l1 = [x for x in l1 if x not in c3]

    c2 = np.random.choice(l2, m2, replace=False)
    l1 = [x for x in l1 if x not in c2]
    c1 = np.random.choice(l1, m1, replace=False)
    idxs_users.append(c1)
    idxs_users.append(c2)
    idxs_users.append(c3)
    for _ in range(m1):
        ration_users.append(0)
    for _ in range(m2):
        ration_users.append(1)
    for _ in range(m3):
        ration_users.append(2)
    return ration_users,idxs_users


def Aggregation_ScaleFL(w, lens, grad_info, global_model_param):
    w_avg = copy.deepcopy(global_model_param)  # largest model
    count = OrderedDict()
    for idx, (k, v) in enumerate(w_avg.items()):  # 遍历所有层，每一层遍历
        parameter_type = k.split('.')[-1]

        count[k] = v.new_zeros(v.size(), dtype=torch.float32)
        tmp_v = v.new_zeros(v.size(), dtype=torch.float32)
        for m in range(len(w)):  # 遍历所有用户
            if grad_info[m][idx]:
                if parameter_type == 'weight':
                    if v.dim() > 1:  # 卷积  和 线性层
                        d1 = w[m][k].shape[0]
                        d2 = w[m][k].shape[1]
                        tmp_v[:d1, :d2] += w[m][k] * lens[m]  # 第m个客户端的 k 层参数
                        count[k][:d1, :d2] += lens[m]
                    else:  # BN层
                        d1 = w[m][k].shape[0]
                        tmp_v[:d1] += w[m][k] * lens[m]
                        count[k][:d1] += lens[m]
                else:
                    d1 = w[m][k].shape[0]
                    tmp_v[:d1] += w[m][k] * lens[m]
                    count[k][:d1] += lens[m]


        tmp_v[count[k] > 0] = tmp_v[count[k] > 0].div_(count[k][count[k] > 0])
        tmp_v[count[k] == 0] = global_model_param[k][count[k] == 0]
        w_avg[k] = tmp_v

    return w_avg